#
# Copyright (C) 2023, Inria
# GRAPHDECO research group, https://team.inria.fr/graphdeco
# All rights reserved.
#
# This software is free for non-commercial, research and evaluation use
# under the terms of the LICENSE.md file.
#
# For inquiries contact  george.drettakis@inria.fr
#

import math
from pathlib import Path
from typing import List, Dict

import torch
from torch import Tensor
from utils.sh_utils import eval_sh
from diff_gaussian_rasterization import GaussianRasterizationSettings, GaussianRasterizer

def render(pc, vp_cam, bg_color, scaling_modifier=1.0) -> Dict[str, Tensor]:

    # Set up rasterization configuration
    raster_settings = GaussianRasterizationSettings(
                    image_height=vp_cam.image_height,
                    image_width=vp_cam.image_width,
                    tanfovx=math.tan(vp_cam.FoVx * 0.5),
                    tanfovy=math.tan(vp_cam.FoVy * 0.5),
                    bg=bg_color,
                    scale_modifier=scaling_modifier,
                    viewmatrix=vp_cam.world_view_transform,
                    projmatrix=vp_cam.full_proj_transform,
                    sh_degree=pc.active_sh_degree,
                    campos=vp_cam.camera_center,
                    prefiltered=False,
                    debug=pc.hp.debug,
                )
    rasterizer = GaussianRasterizer(raster_settings=raster_settings)

    # Create zero tensor. We will use it to make pytorch return gradients of the 2D (screen-space) means
    screenspace_points = torch.zeros_like(pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device='cuda')
    
    try: screenspace_points.retain_grad()
    except: pass
    
    means3D = pc.get_xyz
    means2D = screenspace_points
    opacity = pc.get_opacity
    scales = pc.get_scaling
    rotations = pc.get_rotation    
    
    shs = pc.get_features
    colors_precomp = None
    if pc.precomp_color_from_SHs:
        shs_view = shs.transpose(1, 2).view(-1, 3, (pc.max_sh_degree+1)**2)
        dir_pp = (means3D - vp_cam.camera_center.repeat(shs.shape[0], 1))
        dir_pp_normalized = dir_pp / dir_pp.norm(dim=1, keepdim=True)
        sh2rgb = eval_sh(pc.active_sh_degree, shs_view, dir_pp_normalized)
        colors_precomp = torch.clamp(sh2rgb * 2.0, -pc.color_scale, pc.color_scale)
        shs = None
    
    rendered_image, radii = rasterizer(
                                means3D=means3D,
                                means2D=means2D,
                                shs=shs,
                                colors_precomp=colors_precomp,
                                opacities=opacity,
                                scales=scales,
                                rotations=rotations,
                                cov3D_precomp=None)
                  
    return {
        'render': rendered_image,
        'viewspace_points': screenspace_points,
        'visibility_filter': radii > 0,
        'radii': radii}
